package main

import (
	"bytes"
	"crypto/elliptic"
	"crypto/rand"
	"encoding/asn1"
	"encoding/binary"
	"encoding/gob"
	"encoding/json"
	"encoding/pem"
	"errors"
	"fmt"
	"io"
	"log"
	"math/big"
	"os"

	"github.com/tjfoc/gmsm/sm2"
	"github.com/tjfoc/gmsm/sm3"

	// "github.com/tjfoc/gmsm/sm4"

	"github.com/unidoc/unipdf/v3/common/license"
	"github.com/unidoc/unipdf/v3/model"
	// "github.com/pdfcpu/pdfcpu/pkg/api"
	// "github.com/pdfcpu/pdfcpu/pkg/pdfcpu"
	// "github.com/unidoc/unipdf/v3/common"
	// "github.com/unidoc/unipdf/v3/extractor"
	// "github.com/unidoc/unipdf/v3/model"
)

func systeminit() {

	priv, err := sm2.GenerateKey(rand.Reader) // 生成密钥对
	if err != nil {
		log.Fatal(err)
	}

	privateKey := priv.D.Bytes()
	publicKeyX := priv.PublicKey.X.Bytes()
	publicKeyY := priv.PublicKey.Y.Bytes()
	privateKeyHex := fmt.Sprintf("%x", privateKey)
	publicKeyXHex := fmt.Sprintf("%x", publicKeyX)
	publicKeyYHex := fmt.Sprintf("%x", publicKeyY)
	fmt.Println("Private Key :", privateKeyHex)
	fmt.Println("Public Key X :", publicKeyXHex)
	fmt.Println("Public Key Y :", publicKeyYHex)

	//公开参数写入json
	publicKeyStr := "(" + publicKeyXHex + "," + publicKeyYHex + ")"
	// 创建 Data 结构体对象
	data := Data{
		P_PrimeField:      "fffffffeffffffffffffffffffffffffffffffff00000000ffffffffffffffff",
		GeneratorPoint:    "(32c4ae2c1f1981195f9904466a39c9948fe30bbff2660be1715a4589334c74c7,bc3736a2f4f6779c59bdcee36b692153d0a9877cc62a474002df32e52139f0a0)",
		Order_BasicPoint:  "fffffffeffffffffffffffffffffffff7203df6b21c6052b53bbf40939d54123",
		CurveParameters_a: "fffffffeffffffffffffffffffffffffffffffff00000000fffffffffffffffc",
		CurveParameters_b: "28e9fa9e9d9f5e344d5a9e4bcf6509a7f39789f515ab8f92ddbcbd414d940e93",
		PublicKey_CA:      publicKeyStr,
	}
	jsonData, err := json.MarshalIndent(data, "", "  ")
	if err != nil {
		fmt.Println("Error marshaling JSON:", err)
		return
	}
	// 写入 JSON 文件
	filePath := "parameters.json"
	err = os.WriteFile(filePath, jsonData, 0644)
	if err != nil {
		fmt.Println("Error writing JSON file:", err)
		return
	}

	//保存私钥
	// 创建一个PEM块
	block := &pem.Block{
		Type:  "ECC_KEY_TYPE_HERE", // 例如："RSA PRIVATE KEY"
		Bytes: privateKey,
	}

	// 将PEM块编码为字节切片
	pemBytes := pem.EncodeToMemory(block)

	// 将PEM字节切片写入文件
	err = os.WriteFile("privatekey.pem", pemBytes, 0644)
	if err != nil {
		log.Fatalf("Failed to write data to key.pem: %s", err)
	}
}

func init() {
	// Make sure to load your metered License API key prior to using the library.
	// If you need a key, you can sign up and create a free one at https://cloud.unidoc.io
	// err := license.SetMeteredKey(os.Getenv(`UNIDOC_LICENSE_API_KEY`))
	err := license.SetMeteredKey("3be269bec94ac7b533df95068bb583aac7459adab8582334db078435868af5f5")
	if err != nil {
		panic(err)
	}
	gob.Register(sm2.P256Sm2())

}

func SplitPdf(inputPath string, segmentpos []int) error {
	filePath := "./sourceFile/" + inputPath //[:len(inputPath)-4]
	f, err := os.Open(filePath)
	if err != nil {
		return err
	}

	defer f.Close()

	pdfReader, err := model.NewPdfReader(f)
	if err != nil {
		return err
	}
	i := 1
	var pdfWriter model.PdfWriter
	outputPath := "./access/" + inputPath[:len(inputPath)-4]
	// var filePath string
	for _, j := range segmentpos {
		pdfWriter = model.NewPdfWriter()
		for ; i <= j; i++ {
			page, err := pdfReader.GetPage(i)
			if err != nil {
				return err
			}
			if err := pdfWriter.AddPage(page); err != nil {
				return err
			}
		}
		outputFile := fmt.Sprintf(outputPath+"/%d.pdf", i-1)
		f, err = os.Create(outputFile)
		if err != nil {
			return err
		}
		err = pdfWriter.Write(f)
		f.Close()
		if err != nil {
			return err
		}
		i = 1
	}
	return nil
}

var one = new(big.Int).SetInt64(1)

type resm2Cipher struct {
	XCoordinate *big.Int
	YCoordinate *big.Int
	HASH        []byte
	CipherText  []byte
	Rk          []byte
}

type sm2Cipher struct {
	XCoordinate *big.Int
	YCoordinate *big.Int
	HASH        []byte
	CipherText  []byte
}

func zeroByteSlice() []byte {
	return []byte{
		0, 0, 0, 0,
		0, 0, 0, 0,
		0, 0, 0, 0,
		0, 0, 0, 0,
		0, 0, 0, 0,
		0, 0, 0, 0,
		0, 0, 0, 0,
		0, 0, 0, 0,
	}
}

var (
	C1C3C2 = 0
	C1C2C3 = 1
	En     = 0
	ReEn   = 1
)

type Data struct {
	P_PrimeField      string
	GeneratorPoint    string
	Order_BasicPoint  string
	CurveParameters_a string
	CurveParameters_b string
	PublicKey_CA      string
}

type Client struct {
	ID string
	PublicKey
}

type PublicKey struct {
	elliptic.Curve
	X, Y *big.Int
}

type PrivateKey struct {
	PublicKey
	D *big.Int
}

// 验证用户发来的CSR文件
func verifyCSR(CSRname string, signname string, BPub string) bool {
	//读CSR文件
	CSR, err := os.ReadFile(CSRname)
	if err != nil {
		log.Fatal(err)
	}

	//读sign文件
	signfile, err := os.Open("./sm2Key/sign.dat")
	if err != nil {
		log.Fatal(err)
	}
	defer signfile.Close()

	var sign []byte
	decoder := gob.NewDecoder(signfile)
	err = decoder.Decode(&sign)
	if err != nil {
		log.Fatal("解码签名失败:", err)
	}

	//读CA的曲线
	prepub, err := os.Open(BPub)
	if err != nil {
		log.Fatal("打开文件失败:", err)
	}
	defer prepub.Close()

	// 使用gob解码器还原公钥
	var Bpub *sm2.PublicKey
	decoder = gob.NewDecoder(prepub)
	if err := decoder.Decode(&Bpub); err != nil {
		log.Fatal("解码失败:", err)
	}

	//验证签名
	isok := Bpub.Verify(CSR, sign)
	return isok
}

// 用户对象的加密方法封装, 其中c为文件所有者实例，cli2为被授权者实例
func (c *Client) EncryptAsn1(cli2 *Client, data []byte, random io.Reader, authorization []byte) ([]byte, error) {
	//
	ciphertext, rk, err := Encrypt(&c.PublicKey, &cli2.PublicKey, data, random, C1C3C2, authorization)
	if err != nil {
		return nil, err
	}
	return CipherMarshal(ciphertext, rk, En)
}

// 服务器重加密方法
func ReEncryptAsn1(data []byte) ([]byte, error) {
	cipher, rk, err := CipherUnmarshal_server(data)
	if err != nil {
		return nil, err
	}
	recipher := ReEncrypt(cipher, rk)
	return CipherMarshal(recipher, rk, ReEn)
}

// 用户对象的解密方法封装
func (priv *PrivateKey) DecryptAsn1(data []byte, authorization []byte) ([]byte, error) {
	cipher, err := CipherUnmarshal_client(data)
	if err != nil {
		return nil, err
	}
	return Decrypt(priv, cipher, C1C3C2, authorization)
}

// 将加密密文转换为结构体的形式，并使用asn1编码
func CipherMarshal(data1 []byte, data2 []byte, mode int) ([]byte, error) {
	data1 = data1[1:]
	x := new(big.Int).SetBytes(data1[:32])
	y := new(big.Int).SetBytes(data1[32:64])
	hash := data1[64:96]
	cipherText := data1[96:]
	switch mode {
	case En:
		return asn1.Marshal(resm2Cipher{x, y, hash, cipherText, data2})
	case ReEn:
		return asn1.Marshal(sm2Cipher{x, y, hash, cipherText})
	default:
		return nil, nil
	}
}

// resm2的asn.1编码格式转C1|C3|C2拼接格式和rk
func CipherUnmarshal_server(data []byte) ([]byte, []byte, error) {
	var cipher resm2Cipher
	_, err := asn1.Unmarshal(data, &cipher)
	if err != nil {
		return nil, nil, err
	}
	x := cipher.XCoordinate.Bytes()
	y := cipher.YCoordinate.Bytes()
	hash := cipher.HASH
	if err != nil {
		return nil, nil, err
	}
	cipherText := cipher.CipherText
	if err != nil {
		return nil, nil, err
	}
	rk := cipher.Rk
	if err != nil {
		return nil, nil, err
	}
	if n := len(x); n < 32 {
		x = append(zeroByteSlice()[:32-n], x...)
	}
	if n := len(y); n < 32 {
		y = append(zeroByteSlice()[:32-n], y...)
	}
	c := []byte{}
	c = append(c, x...)          // x分量
	c = append(c, y...)          // y分
	c = append(c, hash...)       // x分量
	c = append(c, cipherText...) // y分
	return append([]byte{0x04}, c...), rk, nil
}

// sm2的asn.1编码格式转C1|C3|C2'拼接格式
func CipherUnmarshal_client(data []byte) ([]byte, error) {
	var cipher sm2Cipher
	_, err := asn1.Unmarshal(data, &cipher)
	if err != nil {
		return nil, err
	}
	x := cipher.XCoordinate.Bytes()
	y := cipher.YCoordinate.Bytes()
	hash := cipher.HASH
	if err != nil {
		return nil, err
	}
	cipherText := cipher.CipherText
	if err != nil {
		return nil, err
	}
	if n := len(x); n < 32 {
		x = append(zeroByteSlice()[:32-n], x...)
	}
	if n := len(y); n < 32 {
		y = append(zeroByteSlice()[:32-n], y...)
	}
	c := []byte{}
	c = append(c, x...)          // x分量
	c = append(c, y...)          // y分
	c = append(c, hash...)       // x分量
	c = append(c, cipherText...) // y分
	return append([]byte{0x04}, c...), nil
}

// 生成随机数k
func randFieldElement(c elliptic.Curve, random io.Reader) (k *big.Int, err error) {
	if random == nil {
		random = rand.Reader //If there is no external trusted random source,please use rand.Reader to instead of it.
	}
	params := c.Params()
	b := make([]byte, params.BitSize/8+8)
	_, err = io.ReadFull(random, b)
	if err != nil {
		return
	}
	k = new(big.Int).SetBytes(b)
	n := new(big.Int).Sub(params.N, one)
	k.Mod(k, n)
	k.Add(k, one)
	return
}

// 整数转字节,大端字节序
func intToBytes(x int) []byte {
	var buf = make([]byte, 4)
	binary.BigEndian.PutUint32(buf, uint32(x))
	return buf
}

// sm2中的密钥派生函数
func kdf(length int, x ...[]byte) ([]byte, bool) {
	var c []byte
	ct := 1
	h := sm3.New()
	for i, j := 0, (length+31)/32; i < j; i++ {
		h.Reset()
		for _, xx := range x {
			h.Write(xx)
		}
		h.Write(intToBytes(ct))
		hash := h.Sum(nil)
		if i+1 == j && length%32 != 0 {
			c = append(c, hash[:length%32]...)
		} else {
			c = append(c, hash...)
		}
		ct++
	}
	for i := 0; i < length; i++ {
		if c[i] != 0 {
			return c, true
		}
	}
	return c, false
}

// 本函数的功能：sm4密钥加密 ，重加密密钥生成, 注意返回字节数据的第一个字节不是密文，用于表示非压缩格式
// data文件所有者加密sm4密钥,  其中pub1是a的公钥（比特数组化后）, pub2是b的公钥, mode use 0, authorization是授权信息
func Encrypt(pub1 *PublicKey, pub2 *PublicKey, data []byte, random io.Reader, mode int, authorization []byte) ([]byte, []byte, error) {
	length := len(data)
	for {
		curve := pub1.Curve
		k, err := randFieldElement(curve, random)
		if err != nil {
			return nil, nil, err
		}
		x1, y1 := curve.ScalarBaseMult(k.Bytes())
		x2, y2 := curve.ScalarMult(pub1.X, pub1.Y, k.Bytes())
		x3, y3 := curve.ScalarMult(pub2.X, pub2.Y, k.Bytes())
		x1Buf := x1.Bytes()
		y1Buf := y1.Bytes()
		x2Buf := x2.Bytes()
		y2Buf := y2.Bytes()
		x3Buf := x3.Bytes()
		y3Buf := y3.Bytes()
		if n := len(x1Buf); n < 32 {
			x1Buf = append(zeroByteSlice()[:32-n], x1Buf...)
		}
		if n := len(y1Buf); n < 32 {
			y1Buf = append(zeroByteSlice()[:32-n], y1Buf...)
		}
		if n := len(x2Buf); n < 32 {
			x2Buf = append(zeroByteSlice()[:32-n], x2Buf...)
		}
		if n := len(y2Buf); n < 32 {
			y2Buf = append(zeroByteSlice()[:32-n], y2Buf...)
		}
		if n := len(x3Buf); n < 32 {
			x3Buf = append(zeroByteSlice()[:32-n], x3Buf...)
		}
		if n := len(y3Buf); n < 32 {
			y3Buf = append(zeroByteSlice()[:32-n], y3Buf...)
		}
		c := []byte{}
		c = append(c, x1Buf...) // C1x分量
		c = append(c, y1Buf...) // C2y分量
		tm := []byte{}
		tm = append(tm, x3Buf...)
		tm = append(tm, data...)
		tm = append(tm, y3Buf...)
		h := sm3.Sm3Sum(tm) //C3
		c = append(c, h...)
		ct, ok := kdf(length, x2Buf, y2Buf) // C2
		if !ok {
			continue
		}
		c = append(c, ct...)
		for i := 0; i < length; i++ {
			c[96+i] ^= data[i]
		}
		rk2, ok := kdf(length, x3Buf, y3Buf, authorization) //这里长度也使用length是由代理服务其中重加密过程决定的
		if !ok {
			continue
		}
		rk := make([]byte, length)
		for i := 0; i < len(ct); i++ {
			rk[i] = ct[i] ^ rk2[i]
		}
		switch mode {

		case C1C3C2:
			return append([]byte{0x04}, c...), rk, nil
		case C1C2C3:
			c1 := make([]byte, 64)
			c2 := make([]byte, len(c)-96)
			c3 := make([]byte, 32)
			copy(c1, c[:64])   //x1,y1
			copy(c3, c[64:96]) //hash
			copy(c2, c[96:])   //密文
			ciphertext := []byte{}
			ciphertext = append(ciphertext, c1...)
			ciphertext = append(ciphertext, c2...)
			ciphertext = append(ciphertext, c3...)
			return append([]byte{0x04}, ciphertext...), rk, nil
		default:
			return append([]byte{0x04}, c...), rk, nil
		}
	}
}

// 重加密函数，放在服务器端，不发送给用户, 用于生成被授权用户可以解密的密文
func ReEncrypt(ciphertext []byte, rk []byte) []byte {
	ciphertext = ciphertext[1:]
	length := len(ciphertext[96:])
	for i := 0; i < length; i++ {
		ciphertext[96+i] = ciphertext[96+i] ^ rk[i]
	}
	return append([]byte{0x04}, ciphertext...)
}

// 被授权者的解密函数, priv为被授权者的私钥
func Decrypt(priv *PrivateKey, data []byte, mode int, authorization []byte) ([]byte, error) {
	switch mode {
	case C1C3C2:
		data = data[1:]
	case C1C2C3:
		data = data[1:]
		c1 := make([]byte, 64)
		c2 := make([]byte, len(data)-96)
		c3 := make([]byte, 32)
		copy(c1, data[:64])             //x1,y1
		copy(c2, data[64:len(data)-32]) //密文
		copy(c3, data[len(data)-32:])   //hash
		c := []byte{}
		c = append(c, c1...)
		c = append(c, c3...)
		c = append(c, c2...)
		data = c
	default:
		data = data[1:]
	}
	length := len(data) - 96
	curve := priv.Curve
	x := new(big.Int).SetBytes(data[:32])
	y := new(big.Int).SetBytes(data[32:64])
	x2, y2 := curve.ScalarMult(x, y, priv.D.Bytes())
	x2Buf := x2.Bytes()
	y2Buf := y2.Bytes()
	if n := len(x2Buf); n < 32 {
		x2Buf = append(zeroByteSlice()[:32-n], x2Buf...)
	}
	if n := len(y2Buf); n < 32 {
		y2Buf = append(zeroByteSlice()[:32-n], y2Buf...)
	}
	c, ok := kdf(length, x2Buf, y2Buf, authorization)
	if !ok {
		return nil, errors.New("Decrypt: failed to derivation")
	}
	for i := 0; i < length; i++ {
		c[i] ^= data[i+96]
	}
	tm := []byte{}
	tm = append(tm, x2Buf...)
	tm = append(tm, c...)
	tm = append(tm, y2Buf...)
	h := sm3.Sm3Sum(tm)
	if !bytes.Equal(h, data[64:96]) {
		return c, errors.New("Decrypt: failed to decrypt")
	}
	return c, nil
}

func main() {

	GUI()

}
